import math

class node():
    def __init__(self,content=None):
        self.content=content
        self.father=-1
        self.children=[]
        self.value=-99
        self.depth=0
        self.visit=0
        self.id=0
        self.labels=[0, 0, 0, 0, 0, 0, 0]   #预留位置
        '''
        labels[0]:是否为完全扩展节点 1:完全扩展 0:未完全扩展
        labels[1]:扩展用哈希值，为多次扩展使用的值
        '''
        
    def add_child(self, tree, content=None):
        child=node(content=content)
        child.id=len(tree)
        child.depth=self.depth+1
        child.father=self.id
        self.children.append(child.id)
        tree.append(child)


MCTS_Iter=int
MCTS_Iter_Init:MCTS_Iter=0
END_OR_NOT=bool
INF=1e5

class Mcts():
    '''
    :root:      添加初始状态
    :clear:     清空蒙特卡洛树所有节点
    :next:      获取当前下一步最优走法
    :select expand simulation backward:
    '''
    
    def __init__(self, c:float, depth_lim:int, end_func:callable, expanded_func:callable, expand_func:callable, value_func:callable):
        self.tree=[]
        self.end_func=end_func              # 返回一个节点是不是终止节点，如果是终止，返回True
        self.expanded_func=expanded_func    # 返回一个节点是不是完全扩展节点,如果完全扩展，返回True
        self.expand_func=expand_func        # 执行扩展函数
        self.value_func=value_func          # 节点估值函数
        self.c=c                            # 探索系数
        self.depth_lim=depth_lim            # 探索深度限制
        self.hashset=set()
        
    def root(self, content):
        root_node=node(content=content)
        root_node.labels[0]=1
        self.tree.append(root_node)
        
    def clear(self):
        self.tree.clear()
        self.hashset.clear()
        

    def select(self):
        '''
        返回第一个 可以扩展的节点 or 叶子节点
        使用循环替代了递归
        
        '''
        iter:MCTS_Iter=MCTS_Iter_Init
        while True:
            # 如果走到了叶子节点，那么直接返回。
            if len(self.tree[iter].children)==0:
                return iter
            # 条件判断，是否为可扩展节点。未完全扩展则返回
            if self.tree[iter].labels[0]==0:
                return iter
            # 选择最有扩展价值的下一个节点
            children=self.tree[iter].children
            iter=max(children, key=lambda x:self.tree[x].value+self.c/math.sqrt(max(self.tree[x].visit, 0.01))) # 避免除0错误
        
    def expand(self, iter:MCTS_Iter)->MCTS_Iter:
        # 设定：一次生成所有可以被接受的子节点方案,对重复采用“鸵鸟算法”
        # 设定：新方案按照可行性降序排列，优先扩展最可行方案
        
        iter_=iter
        while True:
            # 深度限界，停止扩张
            if self.tree[iter_].depth>=self.depth_lim:
                return iter_
            # 如果到达停止状态，停止扩张
            if self.end_func(self.tree[iter_].content)==True:
                return iter_
            # 多次尝试获取接下来状态，如果失败认定为结束节点
            
            
            add_ok=False
            rechoose=3
            for i in range(rechoose):
                next_state=self.expand_func(self.tree[iter_].content, self.tree[iter_].labels[1])
                hash_val=hash(next_state)
                if hash_val in self.hashset:
                    continue
                else:
                    # 添加新状态
                    add_ok=True
                    self.tree[iter_].add_child(self.tree, next_state)
                    self.hashset.add(hash_val)
                    iter_=len(self.tree)-1  # 根据树的规则，添加节点一定添加到了树数组的最后
                    break
            # 判断是否终止
            if not add_ok:
                self.tree[iter_].labels[0]=1
                return iter_
            
        
    def simulation(self, iter:MCTS_Iter):

        # 如果节点已经有估值，那么可以不进行重新估值。估值默认值是-99。
        if self.tree[iter].value==-99:
            value=self.value_func(self.tree[iter].content)
            self.tree[iter].value=value
        #print(self.tree[iter].value)
        return self.tree[iter].value
    
    def backward(self, iter:MCTS_Iter):
        # 不需要递归，循环操作即可
        now_iter=iter
        self.tree[now_iter].visit+=1
        while True:
            father_iter=self.tree[now_iter].father
            father_vis=self.tree[father_iter].visit
            father_value=self.tree[father_iter].value
            now_value=self.tree[now_iter].value
            # 更新价值
            self.tree[father_iter].value=father_value*(float(father_vis)/(father_vis+1.0))+now_value*(1.0/(father_vis+1))
            self.tree[father_iter].visit+=1
            
            # 更新策略
            now_iter=father_iter
            if now_iter<=0:     #root
                break
            
    def next(self):
        '''
        根据最多的访问量确定next
        '''
        root: node = self.tree[0]
        res: node = self.tree[root.children[0]]
        cnt = self.tree[root.children[0]].visit

        # 找到访问量最多的子节点
        for i in root.children:
            if self.tree[i].visit > cnt:
                cnt = self.tree[i].visit  # 修正变量赋值
                res = self.tree[i]

        # 计算改动、差异
        res_board = res.content.board
        root_board = root.content.board
        array = res_board - root_board  # 计算差异
        # 遍历数组，找到第一个差异点
        for i in range(len(array)):  # 从0开始
            for j in range(len(array[0])):  # 从0开始
                if array[i][j] != 0:
                    return [i, j]  # 返回第一个差异点的坐标

        # 如果没有差异，返回None或其他默认值
        return None
            
class state():
    def __init__(self, board, player):
        self.board=board
        self.player=player
        
        
    